import json

import pandas as pd
from fastapi import APIRouter
from starlette.websockets import WebSocket
from client.openai_client import invoke_openai, invoke_laozhang_api
from controller.ai.report import df_to_csv_str
from dao.es_dao.i_iot_run_dao import get_iot_run_data_detail, get_equipment_list
from dao.es_dao.p_health_log_dao import find_health_log_dao

ai_router_equipment = APIRouter(prefix="/test/ai", tags=["人工智能接口"])


@ai_router_equipment.websocket("/equipmentAnalysisWs")
async def websocket_endpoint(ws: WebSocket):
    await ws.accept()
    try:
        while True:
            # 接收前端参数
            raw_data = await ws.receive_text()
            data = json.loads(raw_data)
            # term = {'imei': data['imei']}
            equipment_list = await get_equipment_list(data['term'])
            df_equipment = pd.DataFrame(equipment_list)
            run_log = await get_iot_run_data_detail(data['term'], data['start_time'], data['end_time'])
            df_run = pd.DataFrame(run_log)
            df_run['run_time'] = df_run['intervals'] / 1000
            equipment_list_str = df_to_csv_str(df_equipment)
            run_log_str = df_to_csv_str(df_run)
            query = (f"这是一台或多台保洁公司使用在小区物业的清洁设备的工作数据{run_log_str}，这是设备基础数据{equipment_list_str}，通过字段imei可以关联到设备，"
                     f"开始时间和结束时间是一次运行，说明设备开始工作和结束工作。要求简单详细且高大上，用户是保洁公司的领导，目的是看设备整体工作情况和变化趋势，请忽略单次时长，因为电机就是开开关关走走停停的"
                     f"单次工作时间短，是因为开关一次吸水风机"
                     f"洗地机一般是用在地库里的，扫地机用在社区内外围，路面等。驾驶式的走的快一些。work_amount是每小时的工作效能，单位是work_unit字段，"
                     f"假如你是一名清洁行业和清洁设备行业专家。对设备的运行情况给出分析和建议。"
                     f"你直接面对的是用户，不要说多余的话，因为是展示在页面上的内容"
                     f"要求结果简单明了，不要加入其他的思考内容。")
            model = 'gpt-4o-mini'
            if data.get('model'):
                model = data.get('model')
            await invoke_laozhang_api(query, ws, model=model)
    except Exception as e:
        await ws.close()
        print(f"WebSocket 错误: {str(e)}")
